package com.wanganui.controller;

import com.alibaba.fastjson.JSON;
import com.wanganui.tika.TikaUtil;
import com.wanganui.tika.TikaVo;
import com.wanganui.vector.service.MilvusService;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.QueryResp;
import io.milvus.v2.service.vector.response.SearchResp;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.http.ResponseEntity;
import org.springframework.util.Assert;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.util.*;

/**
 * @author xtwang
 * @des 向量数据库实践
 * @date 2025/2/25 下午5:02
 */
@RestController
@RequestMapping("/api/vector")
@Tag(name = "向量数据库实践", description = "向量数据库实践")
@RequiredArgsConstructor
public class MilvusController {

    private final OllamaEmbeddingModel ollamaEmbeddingModel;
    private final MilvusService milvusService;
    private final TikaUtil tikaUtil;

    @Operation(summary = "获取Embedding")
    @GetMapping("/embedding")
    public ResponseEntity<String> getEmbedding(@RequestParam String text) {
        Assert.notNull(text, "text不能为空");
        float[] embed = ollamaEmbeddingModel.embed(text);
        // 数组全部内容转化成list
        return ResponseEntity.ok(Arrays.toString(embed));
    }

    @Operation(summary = "添加自定义rag数据")
    @GetMapping("/addCustomRagData")
    public ResponseEntity<InsertResp> addCustomRagData(@RequestParam String text) {
        Assert.notNull(text, "text不能为空");
        float[] embed = ollamaEmbeddingModel.embed(text);
        // 插入数据
        InsertResp insert = milvusService.insert(embed, text, JSON.toJSONString(new HashMap<>().put("custom", text)), "custom");
        return ResponseEntity.ok(insert);
    }

    @Operation(summary = "解析文件内容")
    @PostMapping("/extractFileString")
    public ResponseEntity<String> extractFileString(MultipartFile file) {
        return ResponseEntity.ok(tikaUtil.extractTextString(file));
    }

    @Operation(summary = "解析文件内容-LangChina分片")
    @PostMapping("/splitParagraphsLangChain")
    public ResponseEntity<List<String>> splitParagraphsLangChain(MultipartFile file) {
        return ResponseEntity.ok(tikaUtil.splitParagraphsLangChain(tikaUtil.extractTextString(file)));
    }

    @Operation(summary = "解析文件内容-HanLP分片")
    @PostMapping("/splitParagraphsHanLP")
    public ResponseEntity<List<String>> splitParagraphsHanLP(MultipartFile file) {
        return ResponseEntity.ok(tikaUtil.splitParagraphsHanLP(tikaUtil.extractTextString(file)));
    }

    @Operation(summary = "解析文件并分片")
    @PostMapping("/extractFile")
    public ResponseEntity<TikaVo> extractFile(MultipartFile file) {
        return ResponseEntity.ok(tikaUtil.extractText(file));
    }

    @Operation(summary = "上传知识库")
    @PostMapping("/uploadFile")
    public ResponseEntity<InsertResp> uploadFile(MultipartFile file) {
        // 获取文件内容
        TikaVo tikaVo = tikaUtil.extractText(file);
        if (tikaVo != null && Objects.nonNull(tikaVo.getText())) {
            List<float[]> embedList = new ArrayList<>();
            List<String> textList = tikaVo.getText();
            List<String> metadataList = tikaVo.getMetadata();
            List<String> fileNameList = new ArrayList<>();
            for (String s : tikaVo.getText()) {
                embedList.add(ollamaEmbeddingModel.embed(s));
                fileNameList.add(file.getOriginalFilename());
            }
            return ResponseEntity.ok(milvusService.batchInsert(embedList, textList, metadataList, fileNameList));
        }
        return ResponseEntity.ok(null);
    }

    @Operation(summary = "搜索")
    @GetMapping("/search")
    public ResponseEntity<SearchResp> search(@RequestParam String text) {
        Assert.notNull(text, "text不能为空");
        float[] embed = ollamaEmbeddingModel.embed(text);
        // 搜索数据
        return ResponseEntity.ok(milvusService.search(embed));
    }

    @Operation(summary = "分页查询知识库数据集")
    @GetMapping("/query")
    public ResponseEntity<List<QueryResp.QueryResult>> query(@RequestParam Integer page, @RequestParam Integer size, @RequestParam(required = false) String fileName) {
        Assert.notNull(page, "page不能为空");
        Assert.notNull(size, "size不能为空");
        return ResponseEntity.ok(milvusService.query(page, size, fileName));
    }
}
